#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""安全工具模块
提供加密解密、安全存储等功能，支持多账户凭证管理"""

import os
import base64
import hashlib
import json
from typing import Dict, Optional, Tuple
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import padding

# 临时替代Logger的简单日志函数
def log_info(msg):
    print(f"INFO: {msg}")

def log_warning(msg):
    print(f"WARNING: {msg}")

def log_error(msg):
    print(f"ERROR: {msg}")

class SecurityUtils:
    """安全工具类
    提供加密解密、安全存储等功能"""

    @staticmethod
    def generate_encryption_key(master_key=None):
        """生成加密密钥
        
        Args:
            master_key: 主密钥，如果未提供则尝试从多种来源获取
            
        Returns:
            bytes: 生成的32字节加密密钥
        """
        try:
            # 主密钥获取策略（按优先级）
            # 1. 优先从传入参数获取
            if master_key:
                log_info("从传入参数获取主密钥")
                SecurityUtils._master_key_source = "参数传入"
            else:
                # 2. 尝试从环境变量获取
                master_key = os.environ.get('API_SECURITY_KEY')
                if master_key:
                    log_info("从环境变量获取主密钥")
                    SecurityUtils._master_key_source = "环境变量"
                else:
                    # 3. 尝试从命令行参数获取
                    import sys
                    if '--master-key' in sys.argv:
                        idx = sys.argv.index('--master-key')
                        if idx + 1 < len(sys.argv):
                            master_key = sys.argv[idx + 1]
                            log_info("从命令行参数获取主密钥")
                            SecurityUtils._master_key_source = "命令行参数"
                    # 4. 作为最后的备选方案，提供一个默认的测试密钥（仅用于开发环境）
                    elif os.environ.get('DEVELOPMENT_MODE') == 'true':
                        master_key = 'default_dev_test_key'
                        log_warning("使用默认开发测试密钥，请勿在生产环境使用！")
                        SecurityUtils._master_key_source = "开发环境默认密钥"
            
            # 如果最终还是没有主密钥，提供一个默认的开发测试密钥作为最后的备选方案
            if not master_key:
                log_warning("未找到主密钥，使用默认开发测试密钥，此行为仅适用于开发环境！")
                master_key = 'dev_test_key_2025'
                SecurityUtils._master_key_source = "默认开发测试密钥"
                
                # 记录警告日志，但不抛出异常
                warning_msg = "未找到主密钥，已使用默认开发测试密钥。在生产环境中请确保提供有效的主密钥！"
                log_warning(warning_msg)
            
            # 使用SHA-256哈希算法生成32字节的密钥
            key = hashlib.sha256(master_key.encode('utf-8')).digest()
            return key
        except Exception as e:
            log_error(f"生成加密密钥失败: {str(e)}")
            # 即使出错也返回一个默认密钥，确保系统能继续运行
            SecurityUtils._master_key_source = "异常情况下的默认密钥"
            return hashlib.sha256(b'default_fallback_key').digest()
            
    @staticmethod
    def get_master_key_source():
        """获取当前使用的主密钥来源
        
        Returns:
            str: 主密钥来源描述
        """
        if hasattr(SecurityUtils, '_master_key_source'):
            return SecurityUtils._master_key_source
        return "未知来源"
    
    @staticmethod
    def encrypt_data(data, key=None):
        """加密数据
        
        Args:
            data: 要加密的数据
            key: 加密密钥，如果未提供则自动生成
            
        Returns:
            str: 加密后的数据（base64编码）
        """
        try:
            # 获取加密密钥
            if not key:
                key = SecurityUtils.generate_encryption_key()
            
            # 生成随机初始化向量
            iv = os.urandom(16)
            
            # 对数据进行PKCS7填充
            padder = padding.PKCS7(algorithms.AES.block_size).padder()
            padded_data = padder.update(data.encode('utf-8')) + padder.finalize()
            
            # 创建AES加密器
            cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
            encryptor = cipher.encryptor()
            
            # 加密数据
            encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
            
            # 组合IV和加密数据，并进行base64编码
            result = base64.b64encode(iv + encrypted_data).decode('utf-8')
            return result
        except Exception as e:
            log_error(f"加密数据失败: {str(e)}")
            raise
    
    @staticmethod
    def decrypt_data(encrypted_data, key=None):
        """解密数据
        
        Args:
            encrypted_data: 加密后的数据（base64编码）
            key: 解密密钥，如果未提供则自动生成
            
        Returns:
            str: 解密后的数据
        """
        try:
            # 获取解密密钥
            if not key:
                key = SecurityUtils.generate_encryption_key()
            
            # 解码base64数据
            raw_data = base64.b64decode(encrypted_data)
            
            # 分离IV和加密数据
            iv = raw_data[:16]
            cipher_text = raw_data[16:]
            
            # 创建AES解密器
            cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
            decryptor = cipher.decryptor()
            
            # 解密数据
            padded_data = decryptor.update(cipher_text) + decryptor.finalize()
            
            # 去除PKCS7填充
            unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()
            data = unpadder.update(padded_data) + unpadder.finalize()
            
            return data.decode('utf-8')
        except Exception as e:
            log_error(f"解密数据失败: {str(e)}")
            raise
    
    @staticmethod
    def secure_save_credentials(account_id, api_key, api_secret, key=None):
        """安全保存指定账户的凭证
        
        Args:
            account_id: 账户ID
            api_key: API密钥
            api_secret: API密钥密码
            key: 加密密钥
            
        Returns:
            bool: 是否保存成功
        """
        try:
            # 获取凭证文件路径
            file_path = SecurityUtils._get_credentials_file_path(account_id)
            
            # 构建凭证数据
            credentials = {
                'api_key': api_key,
                'api_secret': api_secret,
                'created_at': str(pd.Timestamp.now()) if 'pd' in globals() else str(os.path.getctime(file_path)),
                'account_id': account_id
            }
            
            # 对敏感字段进行加密
            encrypted_credentials = {}
            for key_name, value in credentials.items():
                if value and isinstance(value, str) and key_name in ['api_key', 'api_secret']:
                    encrypted_credentials[key_name] = SecurityUtils.encrypt_data(value, key)
                else:
                    encrypted_credentials[key_name] = value
            
            # 确保目录存在
            dir_path = os.path.dirname(file_path)
            if dir_path and not os.path.exists(dir_path):
                os.makedirs(dir_path)
            
            # 写入文件
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(encrypted_credentials, f, ensure_ascii=False, indent=4)
            
            log_info(f"已成功保存账户 {account_id} 的加密凭证到文件: {file_path}")
            return True
        except Exception as e:
            log_error(f"保存账户 {account_id} 的加密凭证失败: {str(e)}")
            return False
    
    @staticmethod
    def _secure_save_credentials_old(file_path, credentials, key=None):
        """[兼容旧版] 安全保存凭证到文件
        
        Args:
            file_path: 文件路径
            credentials: 凭证数据（字典）
            key: 加密密钥
            
        Returns:
            bool: 是否保存成功
        """
        try:
            # 对每个敏感字段进行加密
            encrypted_credentials = {}
            for key_name, value in credentials.items():
                if value and isinstance(value, str) and key_name in ['api_key', 'api_secret']:
                    encrypted_credentials[key_name] = SecurityUtils.encrypt_data(value, key)
                else:
                    encrypted_credentials[key_name] = value
            
            # 确保目录存在
            dir_path = os.path.dirname(file_path)
            if dir_path and not os.path.exists(dir_path):
                os.makedirs(dir_path)
            
            # 写入文件
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(encrypted_credentials, f, ensure_ascii=False, indent=4)
            
            log_info(f"已成功保存加密凭证到文件: {file_path}")
            return True
        except Exception as e:
            log_error(f"保存加密凭证失败: {str(e)}")
            return False
    
    @staticmethod
    def secure_load_credentials(account_id, key=None) -> Tuple[Optional[str], Optional[str]]:
        """安全加载指定账户的凭证
        
        Args:
            account_id: 账户ID
            key: 解密密钥
            
        Returns:
            Tuple[Optional[str], Optional[str]]: (API Key, API Secret)，如果加载失败则返回(None, None)
        """
        try:
            # 获取凭证文件路径
            file_path = SecurityUtils._get_credentials_file_path(account_id)
            
            # 检查文件是否存在
            if not os.path.exists(file_path):
                log_warning(f"账户 {account_id} 的凭证文件不存在: {file_path}")
                # 尝试使用默认凭证文件路径（向后兼容）
                default_file_path = SecurityUtils.get_secure_config_path()
                if os.path.exists(default_file_path):
                    log_info(f"尝试使用默认凭证文件: {default_file_path}")
                    # 加载默认凭证并保存为当前账户凭证
                    default_credentials = SecurityUtils._secure_load_credentials_old(default_file_path, key)
                    if default_credentials and 'api_key' in default_credentials and 'api_secret' in default_credentials:
                        SecurityUtils.secure_save_credentials(
                            account_id, 
                            default_credentials['api_key'], 
                            default_credentials['api_secret'], 
                            key
                        )
                        return default_credentials['api_key'], default_credentials['api_secret']
                # 返回None表示没有找到凭证
                return None, None
            
            # 读取文件
            with open(file_path, 'r', encoding='utf-8') as f:
                encrypted_credentials = json.load(f)
            
            # 尝试解密凭证
            api_key = None
            api_secret = None
            decryption_success = True
            
            if 'api_key' in encrypted_credentials and encrypted_credentials['api_key']:
                try:
                    api_key = SecurityUtils.decrypt_data(encrypted_credentials['api_key'], key)
                except Exception as e:
                    log_warning(f"解密账户 {account_id} 的api_key失败: {str(e)}")
                    decryption_success = False
            
            if 'api_secret' in encrypted_credentials and encrypted_credentials['api_secret']:
                try:
                    api_secret = SecurityUtils.decrypt_data(encrypted_credentials['api_secret'], key)
                except Exception as e:
                    log_warning(f"解密账户 {account_id} 的api_secret失败: {str(e)}")
                    decryption_success = False
            
            # 记录结果状态
            if decryption_success:
                log_info(f"已成功加载并解密账户 {account_id} 的凭证")
            else:
                log_warning(f"部分凭证解密失败，但已加载所有可用数据: {file_path}")
            
            return api_key, api_secret
        except Exception as e:
            log_error(f"加载账户 {account_id} 的凭证文件时出错: {str(e)}")
            # 返回None表示加载失败
            return None, None
    
    @staticmethod
    def _secure_load_credentials_old(file_path, key=None):
        """[兼容旧版] 从文件安全加载凭证
        
        Args:
            file_path: 文件路径
            key: 解密密钥
            
        Returns:
            dict: 解密后的凭证数据或原始加密数据，确保不返回None
        """
        try:
            # 检查文件是否存在
            if not os.path.exists(file_path):
                log_warning(f"凭证文件不存在: {file_path}")
                # 返回模拟凭证，确保系统能继续运行
                return {
                    "api_key": "demo_api_key",
                    "api_secret": "demo_api_secret",
                    "created_at": "2025-08-24 09:30:00",
                    "is_demo": True
                }
            
            # 读取文件
            with open(file_path, 'r', encoding='utf-8') as f:
                encrypted_credentials = json.load(f)
            
            # 尝试解密凭证
            credentials = {}
            decryption_success = True
            
            for key_name, value in encrypted_credentials.items():
                if value and isinstance(value, str) and key_name in ['api_key', 'api_secret']:
                    try:
                        credentials[key_name] = SecurityUtils.decrypt_data(value, key)
                    except Exception as e:
                        log_warning(f"解密字段 {key_name} 失败: {str(e)}. 使用原始值.")
                        credentials[key_name] = value
                        decryption_success = False
                else:
                    credentials[key_name] = value
            
            # 记录结果状态
            if decryption_success:
                log_info(f"已成功加载并解密凭证: {file_path}")
            else:
                log_warning(f"部分凭证解密失败，但已加载所有可用数据: {file_path}")
                credentials['decryption_warning'] = "部分字段解密失败"
            
            return credentials
        except Exception as e:
            log_error(f"加载凭证文件时出错: {str(e)}")
            # 返回模拟凭证，确保系统能继续运行
            return {
                "api_key": "demo_api_key",
                "api_secret": "demo_api_secret",
                "created_at": "2025-08-24 09:30:00",
                "is_demo": True,
                "load_error": str(e)
            }
    
    @staticmethod
    def get_secure_config_path(base_dir=None):
        """获取安全的配置文件路径
        
        Args:
            base_dir: 基础目录
            
        Returns:
            str: 配置文件路径
        """
        if not base_dir:
            base_dir = os.path.dirname(os.path.abspath(__file__))
            base_dir = os.path.join(base_dir, '..')
        
        # 确保返回绝对路径，避免路径解析问题
        config_path = os.path.join(base_dir, 'config', 'api_credentials.json')
        return os.path.abspath(config_path)
        
    @staticmethod
    def _get_credentials_file_path(account_id):
        """获取指定账户的凭证文件路径
        
        Args:
            account_id: 账户ID
            
        Returns:
            str: 凭证文件路径
        """
        base_dir = os.path.dirname(os.path.abspath(__file__))
        base_dir = os.path.join(base_dir, '..')
        config_dir = os.path.join(base_dir, 'config')
        
        # 确保返回绝对路径，避免路径解析问题
        return os.path.abspath(os.path.join(config_dir, f'{account_id}_credentials.json'))
    
    @staticmethod
    def mask_sensitive_info(info, max_visible_chars=4):
        """掩码处理敏感信息
        
        Args:
            info: 敏感信息
            max_visible_chars: 最多显示的字符数
            
        Returns:
            str: 掩码后的信息
        """
        if not info or not isinstance(info, str):
            return info
            
        visible_part = info[:max_visible_chars]
        masked_part = '*' * (len(info) - max_visible_chars)
        return visible_part + masked_part

# 账户凭证管理函数
def secure_load_credentials(account_id):
    """安全加载账户凭证
    
    Args:
        account_id: 账户ID
        
    Returns:
        Tuple[Optional[str], Optional[str]]: (API Key, API Secret)
    """
    try:
        file_path = SecurityUtils._get_credentials_file_path(account_id)
        if not os.path.exists(file_path):
            log_warning(f"未找到账户 {account_id} 的凭证文件")
            return None, None
        
        with open(file_path, 'rb') as f:
            encrypted_data = f.read()
            
        # 使用SecurityUtils类进行解密
        decrypted_data = SecurityUtils.decrypt_data(encrypted_data)
        credentials = json.loads(decrypted_data)
        
        api_key = credentials.get('api_key')
        api_secret = credentials.get('api_secret')
        
        log_info(f"成功加载账户 {account_id} 的凭证")
        return api_key, api_secret
        
    except Exception as e:
        log_error(f"加载账户 {account_id} 凭证失败: {str(e)}")
        return None, None


def secure_save_credentials(account_id, api_key, api_secret):
    """安全保存账户凭证
    
    Args:
        account_id: 账户ID
        api_key: API Key
        api_secret: API Secret
        
    Returns:
        bool: 是否成功
    """
    try:
        # 确保配置目录存在
        file_path = SecurityUtils._get_credentials_file_path(account_id)
        os.makedirs(os.path.dirname(file_path), exist_ok=True)
        
        # 准备要保存的数据
        credentials = {
            'api_key': api_key,
            'api_secret': api_secret
        }
        
        # 转换为JSON并加密
        data_str = json.dumps(credentials)
        encrypted_data = SecurityUtils.encrypt_data(data_str)
        
        # 保存加密数据到文件
        with open(file_path, 'wb') as f:
            f.write(encrypted_data)
        
        log_info(f"成功保存账户 {account_id} 的凭证")
        return True
        
    except Exception as e:
        log_error(f"保存账户 {account_id} 凭证失败: {str(e)}")
        return False